from matplotlib import pyplot as plt
import pandas as pd
import statsmodels.formula.api as smf
from sqlalchemy import create_engine
import pymysql
import statsmodels.api as sm
import numpy as np
db_config = {
    'host': '111.231.14.211',
    'user': 'tushare',
    'password': 'root',
    'database': 'tushare',
    'port': 13307,
    'charset':'utf8mb4'
}
engine = create_engine(
    f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}"
)
conn = pymysql.connect(**db_config)
chunk_size = 10000


df = pd.read_sql_query(
    """
    SELECT d.*, m.net_mf_vol , m.sell_elg_vol , m.buy_elg_vol, m.sell_lg_vol, m.buy_lg_vol , i.closes as i_closes, i.vol as i_vol
    FROM date_1 d
    JOIN moneyflows m
    ON d.ts_code = m.ts_code AND d.trade_date = m.trade_date
    LFFT JOIN index_daily i ON d.trade_date = i.trade_date and i.ts_code='000001.SH'
    WHERE d.trade_date BETWEEN '2023-01-01' AND '2023-12-31' AND d.ts_code='000001.SZ'
    """,
    conn,
    chunksize=chunk_size
    ) 
df1 = pd.concat(df, ignore_index=True)

#df1.head()

df1['zd_close'] =df1['closes'].shift(1)

df1['hz_close'] = round((df1['i_closes'] - df1['i_closes'].shift(1)) / df1['i_closes'].shift(1), 2)
df1['hz_vol'] = round((df1['i_vol'].shift(1)- df1['i_vol'].shift(2)) / df1['i_vol'].shift(2), 2)

df1 = df1.dropna(subset=['zd_close']).reset_index(drop=True)


numeric_cols = df1.select_dtypes(include=['number']).columns.tolist()



x =df1[['amount', 'net_mf_vol','sell_elg_vol','buy_lg_vol','hz_close','hz_vol']]


eigenvalues, eigenvectors = np.linalg.eig(np.cov(x, rowvar=False))
print('累计贡献率为: ',round(eigenvalues[:5].sum()/eigenvalues.sum(),4)*100,'%')

n_components = 5
top_eigenvectors = eigenvectors[:, :n_components]


principal_components = np.dot(top_eigenvectors)


principal_components = np.dot(x, top_eigenvectors)
data_pca = pd.concat([df1,pd.DataFrame(principal_components,
                   columns=[f'PC{i+1}' for i in range(n_components)])], axis=1)


X_pca = data_pca[[f'PC{i+1}' for i in range(n_components)]].copy()
X_pca = sm.add_constant(X_pca)


y = df1['zd_close'].copy()


model_pca = sm.OLS(y, X_pca)


result_pca = model_pca.fit()


print("\n回归模型结果:")
print(result_pca.summary())


X_pca_selected = data_pca[[ 'PC1', 'PC3', 'PC4']]

X_pca_selected.columns = ['PC1', 'PC3', 'PC4']

X_pca_selected = sm.add_constant(X_pca_selected)


# y = df1['zd_close].copy()


model_pca_selected = sm.OLS(y, X_pca_selected)


result_pca_selected = model_pca_selected.fit()



print("\n回归模型结果:")
print(result_pca_selected.summary())

X_pca_selected.columns = [['PC1', 'PC3', 'PC4']]

X_pca_selected.columns = ['PC1', 'PC3', 'PC4']

fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(15, 5))

for i, col in enumerate(X_pca_selected.columns):
    axes[i].scatter(X_pca_selected[col], y, s=50, alpha=0.7)
    axes[i].set_xlabe(col)
    axes[i].set_ylabel('(y)')
    axes[i].set_title(f'{col}')
    
    plt.tight_layout()
    plt.show()
    
    for k in range(0.5):
        string_y = f'PC{k+1} = '
        i = eigenvectors[k]
        for j in range(len(i)):
            if i[j] > 0 :
                string_y = string_y + f'+{round(i[j],2)}*X_(j+1)'
            else:
                string_y = string_y + f'+{round(i[j],2)}*X_(j+1)'
        if k!=2 and k!=4:
            print(string_y)